#pragma once

#include <unordered_set>

#include "DBManager/DBAPI.h"
#include "DBManager/Repository/HDF5/Base/SingleRepositoryBase.h"

SG_DBMANAGER_HDF5_NAMESPACE_OPEN

template <typename T>
class MpiRepositoryBase : public SingleRepositoryBase<T>
{
  public:
    using CacheRepositoryBase<T>::GetDatabase;
    explicit MpiRepositoryBase (HDF5Database& database) : SingleRepositoryBase<T> (database) {};

    ~MpiRepositoryBase () {};

#if defined(SG_USE_MPI)
    using MainlyRepositoryBase<T>::ElementIdFilter;
    /**
     * @brief    过滤函数，获取当前分区内的id
     *
     * @param    ids
     * @return
     * @author   wanglong
     * @date     2025.08.06
     */
    virtual void IdFilter (std::vector<SG::DataStructure::Common::Id>& ids) override
    {
        ElementIdFilter (GetDatabase (), ids);  // 调用 ElementIdFilter 函数进行 ID 过滤
    }
#endif

    /**
     * @brief 获取表单的行数（数据总量）。
     *
     * @return 表单的行数。
     */
    virtual std::size_t Size () override
    {
        if (GetDatabase ().GetPartitionisFilter ())
        {
            auto* allId =
                this->m_cacheHelper.template GetVectorFast<SG::DataStructure::Common::Id> (this->GetMainDatasetInfo ().GetPath ());
            if (allId == nullptr)
            {
                allId  = new std::vector<SG::DataStructure::Common::Id> ();
                *allId = this->GetAllId ();
                this->m_cacheHelper.template CacheVector (this->GetMainDatasetInfo ().GetPath (), allId);
            }

            return allId->size ();  ///< 根据 ID 查找实体
        }
        else
        {
            return SingleRepositoryBase<T>::Size ();  ///< 调用基类方法查找实体
        }
    }

    /**
     * @brief 向缓存中添加多个实体。当缓存达到最大数量时，将调用 `Dump()` 方法进行处理。
     *
     * @param entity 要添加的实体。
     */
    SG::DataStructure::Common::Status Add (std::vector<T>& entities) override
    {
        this->m_cacheHelper.Remove (this->GetMainDatasetInfo ().GetPath ());
        return SingleRepositoryBase<T>::Add (entities);
    }

    /**
     * @brief 向缓存中添加一个实体。当缓存达到最大数量时，将调用 `Dump()` 方法进行处理。
     *
     * @param entity 要添加的实体。
     */
    SG::DataStructure::Common::Status Add (T& entity) override
    {
        this->m_cacheHelper.Remove (this->GetMainDatasetInfo ().GetPath ());
        return SingleRepositoryBase<T>::Add (entity);
    }

    /**
     * @brief 根据索引查找实体。
     *
     * @param index 要查找的索引。
     * @param _OUT entity 查找到的实体将存储在此参数中。
     * @return 如果索引有效且实体成功加载，则返回 `true`；否则返回 `false`。
     *
     * @throw std::runtime_error 如果索引超出范围。
     */
    virtual SG::DataStructure::Common::Status FindByIndex (::std::size_t index, T& _OUT entity)
    {
        if (GetDatabase ().GetPartitionisFilter ())
        {
            auto* allId =
                this->m_cacheHelper.template GetVectorFast<SG::DataStructure::Common::Id> (this->GetMainDatasetInfo ().GetPath ());
            if (allId == nullptr)
            {
                allId  = new std::vector<SG::DataStructure::Common::Id> ();
                *allId = this->GetAllId ();
                this->m_cacheHelper.template CacheVector (this->GetMainDatasetInfo ().GetPath (), allId);
            }

            return this->FindById (allId->at (index), entity);  ///< 根据 ID 查找实体
        }
        else
        {
            return SingleRepositoryBase<T>::FindByIndex (index, entity);  ///< 调用基类方法查找实体
        }
    }
};

SG_DBMANAGER_HDF5_NAMESPACE_CLOSE